/**
  ETFAna project, Anyang Normal University && IMP-CAS
  \class ETFTracker
  \brief a track finder, consisting of DCs, tracking algorithm, and tracks
  \author SUN Yazhou, asia.rabbit@163.com
  \since 2021-11-01
  \date 2024-03-07 last modified
  \attention
  changelog
  <table>
  <tr>  <th>Date         <th>Author      <th>Description                   </tr>
  <tr>  <td>2021-11-01   <td>Asia Sun    <td>file created                  </tr>
  </table>

  \copyright Copyright (c) 2021-2024 Anyang Normal U. && IMP-CAS with LGPLv3 LICENSE
*/

#include <iostream>
#include <string>
#include <TClonesArray.h>
#include <TClass.h>
#include <TTree.h>
#include "ETFRun.h"
#include "ETFTracker.h"
#include "ETFDCHit.h"
#include "ETFVTrack.h"
#include "ETF3DTrack.h"
#include "ETFTrackTa.h"
#include "ETFTrackTa4.h"
#include "ETFMsg.h"
#include "ETFVisual.h"
#include "ETFDetectorManager.h"
#include "ETFParaManager.h"
#include "ETFTimer.h"

using std::cout;
using std::endl;
using std::string;

#define er ETFMsg::Error

ETFTracker::ETFTracker(const char *name) : fBroken(0), fTracks{0}, f3DTracks(0),
fTracksOut{0}, fHasMapped{0}, fNLayer(0), fdm{nullptr}, fpm{nullptr}, fcfg(0),
fTimer{0}, D2THREPD(-999.), fChiPD(-999.), fMultiThrePerL(-1), fDebug(false){}

ETFTracker::ETFTracker(const char *name, const vector<string> &dcs, const string &t0,
const string &t1) : fBroken(0), fTracks{0}, f3DTracks(0), fTracksOut{0}, fHasMapped{0},
fNLayer(0), fdm{nullptr}, fpm{nullptr}, fcfg(0), fTimer{0}, D2THREPD(-999.),
fChiPD(-999.), fMultiThrePerL(-1), fDebug(false){
  if(!strcmp(name, "")) er("ETFTracker", "ctor: empty name");
  fdm = ETFDetectorManager::Instance();
  fpm = ETFParaManager::Instance();
  const ETFYaml &cc = (*(fcfg = fpm->GetMisPara()))[name];
  // add the tracking dcs to the dc vector
  for(const string &dc : dcs){
    auto d = dynamic_cast<ETFVDC *>(fdm->GetDetector(dc));
    if(d) fDCs.push_back(d);
  } // end for over dcs
  if(fDCs.size() <= 1){ fBroken = true; return; } // no need to continue
  // add t0 detectors
  if("" != t0) fTimer[0] = dynamic_cast<ETFTimer *>(fdm->GetDetector(t0));
  if("" != t1) fTimer[1] = dynamic_cast<ETFTimer *>(fdm->GetDetector(t1));
  // read the tracking control parameters
  if(cc["D2ThrePD"]) D2THREPD = cc["D2ThrePD"].as<double>();
  if(cc["chiPD"]) fChiPD = cc["chiPD"].as<double>();
  fMultiThrePerL = 4;
  if(cc["multiThrePerL"]) fMultiThrePerL = cc["multiThrePerL"].as<int>();
} // end ctor

ETFTracker::~ETFTracker(){
  for(auto &p : fTracks) if(p) delete p;
  for(auto &p : fTracksOut) if(p) delete p;
  if(f3DTracks) delete f3DTracks;
} // end dtor

void ETFTracker::Initialize(){
  if(fBroken) return;
  fHits.clear();
  for(auto &p : fTracks) if(p) p->Clear();
  for(auto &p : fTracksOut) if(p) p->Clear();
  if(f3DTracks) f3DTracks->Clear();
  memset(fHasMapped, 0, sizeof(fHasMapped));
  fNLayer = 0;
} // end member function Initialize

// register det to main tree by name
void ETFTracker::Register(){
  if(fBroken) return;
  TTree *tree = ETFRun::Instance()->GetTree();
  for(auto &t : fTracksOut) if(t){
    // branchname, foldername, array address, persistent or transient
    tree->Branch(t->GetName(), t);
  } // end if
  if(f3DTracks) tree->Branch(f3DTracks->GetName(), f3DTracks);
} // end member function Register

// return if string s contains string p
inline bool has(const string &s, const char *p){ return s.find(p) != string::npos; }
// remove nullptr in TClonesArray to output clones array - IMPORTANT
void ETFTracker::Clean(TClonesArray *cIn, TClonesArray *cOut){
  cOut->Clear();
  int i = 0;
  for(const auto &p : *cIn){ // this scoped for skips nullptrs automatically
    ETFVTrack *t = static_cast<ETFVTrack *>(cOut->ConstructedAt(i++));
    t->copy(static_cast<ETFVTrack *>(p));
    if(has(GetName(), "lobal")) t->SetHitFiredStatus(3);
    else t->SetHitFiredStatus(2); // mark to be distinctive from global trk hits
    // update drift time in case that the hits are altered by
    // other (defeated) tracks with different TOFW strips
    if(t->GetFitStatus() != FitStatus::kUNFITTED){
      t->UpdateDriftTime();
      if(!SetDriftDistance(t) || !t->Fit()) er(GetName(), "Clean: Shit happens");
    } // end if
    t->UpdateChi(); // this is the final chi
  } // end for

  // set distinctive colors for preTaTrk part and postTaTrk part of a ETFTrackTa4 trk
  for(const auto &p : *cIn){
    ETFTrackTa4 *t = dynamic_cast<ETFTrackTa4 *>(p);
    if(t){ // an ETFTrackTa4 track
      t->GetPreTaTrack()->SetLineColor(kRed);
      t->GetPostTaTrack()->SetLineColor(kBlue);
    } // end if
  } // end for
} // end member function Clean

// update the hit position at det
void ETFTracker::UpdateHitPos(ETFDetector *d){
  if(!d) return;
  for(auto &c : fTracksOut) if(c) for(const auto &t : *c) if(t)
    static_cast<ETFVTrack *>(t)->UpdateHitPos(d);
  if(f3DTracks) for(const auto &t : *f3DTracks) if(t)
    static_cast<ETF3DTrack *>(t)->UpdateHitPos(d);
} // end member function UpdateHitPos
// update the hit position of the tracks in the tracker
void ETFTracker::UpdateHitPos(){
  for(auto &d : fDCs) if(d) UpdateHitPos(d);
} // end member function UpdateHitPos

// construct an object in the first empty slot in c
TObject *ETFTracker::ConstructedAt(TClonesArray *c){
  int i = 0;
  while(i < 1000){
    if(!(c->At(i))) return c->ConstructedAt(i);
    i++;
  } // end while
  // ETFMsg::Warn(GetName(), "ConstructedAt: Odd... Should never reach here");
  return nullptr;
} // end member function ConstructedAt

void ETFTracker::SetMapped(bool opt, TrackType type){
  fHasMapped[short(type)] = opt;
  if(!opt) fNLayer = 0; // to initialize map so that it starts from recursion depth 0
} // end member function SetMapped

int ETFTracker::GetNofTracks(TrackType type) const{
  if(!GetTracks(type)) return 0;
  return GetTracks(type)->GetEntriesFast();
} // end member function GetNofTracks

int ETFTracker::GetNof3DTracks() const{
  return Get3DTracks()->GetEntriesFast();
} // end member function GetNofTracks

// the map method is just for exhausting all the possible hit combinations
// using recursion to achieve embedded loops
void ETFTracker::Map(TrackType type){
  // sense wire layer id in the DC array
  const int NLAYER = fDCs.size()*2; // number of sense wire layers in the DC array
  // fNLayer = 0: recursion depth 0
  if(!fNLayer){
    if(!fHasMapped[int(type)]) fHasMapped[int(type)] = true; // first calling of Map
    else return;
  } // end outer if

  TClonesArray *layer = fDCs[fNLayer/2]->GetHitsArray(type, fNLayer%2);
  int n = layer->GetEntriesFast();
  // if too many multiplicity, nullify the whole layer
  if(!GoodLayer(fDCs[fNLayer/2], type, fNLayer%2)) n = 0;
  for(int i = 0; i <= n; i++){
    if(i != n){
      ETFDCHit *h = static_cast<ETFDCHit *>(layer->At(i));
      if(h->GetFiredStatus() == 0) continue; // a nullified hit
      fHits[fNLayer] = h;
      fHits[fNLayer]->SetLayerId(fNLayer);
    } // end if
    else fHits[fNLayer] = nullptr; // include the possibility that all hits are invalid in the layer
    // the last layer reached, now we can do the tracking
    if(fNLayer + 1 == NLAYER) TestCombination(type); // test if the combination'd be accepted
    else{ fNLayer++; Map(type); } // recursive calling
  } // end for over i
  if(!fNLayer) Clean(fTracks[short(type)], fTracksOut[short(type)]); // remove nullptrs
  if(fNLayer > 0) fNLayer--; // fNLayer == 0: the outermost layer of Map, need not decrement
} // end member function Map

bool ETFTracker::Compare(ETFVTrack *trk, TrackType type){
  // no conclusions could be drawn for U or V track projections
  if(TrackType::kX != type && TrackType::kY != type)
    er("ETFTracker", "Compare: supposed only to be called for X or Y tracks");
  // compare the new track with each of the tracks already stored
  for(auto oldtrk : *fTracks[short(type)]){
    ETFVTrack *trk0 = static_cast<ETFVTrack *>(oldtrk);
    if(trk == trk0) continue; // skip the same tracks
    if(!Compare(trk, trk0, type)) return false;
  } // end for over tracks
  return true;
} // end member function Compare

void ETFTracker::TrackMerge(){
  ETFMsg::Warn("ETFTracker", "TrackMerge: Please implement this in your derived class. thx");
} // end member function TrackMerge

// compare trk with all the stored tracks of the same type
// if trk is considered to belong to a different particle
// it would then be accepted and added to the list of tracks
// return true if the new track is accepted
// \param trk1: newtrk; trk0: oldtrk
bool ETFTracker::Compare(ETFVTrack *trk1, ETFVTrack *trk0, TrackType type){
  // TODO: user input configuration expected here
  static const int VICINITY = 3; // min nof cells 2 seperate tracks can be apart from for
  const int NL = fDCs.size() * 2; // number of wire layers
  // number of points (hits) in the new and old tracks
  const int np0 = trk0->GetNofHits(), np1 = trk1->GetNofHits();
  int vicinityViaCnt = 0; // count of vicinity vialations
  int nCom = 0; // nof layers where trk0 and trk1 both have hits
  for(int i = 0; i < NL; i++) if((*trk0)[i] && (*trk1)[i]) nCom++;
  // -------------- the new track has less hits than the old track -------------//
  if(np1 < np0){
    for(int i = 0; i < NL; i++)
      if((*trk0)[i] && (*trk1)[i] && abs((*trk0)[i]->GetNu()-(*trk1)[i]->GetNu()) > VICINITY)
        // if(++vicinityViaCnt == np1) return true; // the two tracks are different
        if(++vicinityViaCnt == nCom) return true; // the two tracks are different
    return false; // the 2 tracks are too close: newtrk is considered part of the old track
  } // end if(np1 < np0)
  // -------------- the new track has more hits than the old track -------------//
  if(np1 > np0){
    for(int i = 0; i < NL; i++)
      if((*trk0)[i] && (*trk1)[i] && abs((*trk0)[i]->GetNu()-(*trk1)[i]->GetNu()) > VICINITY)
        // if(++vicinityViaCnt == np0) return true; // the two tracks are different
        if(++vicinityViaCnt == nCom) return true; // the two tracks are different
    // the 2 tracks are too close: newtrk beats the old track - oldtrk is dropped
    fTracks[short(type)]->Remove(trk0);
    return true;
  } // end if(np1 > np0)
  //-------------- the new track has equal hits to the old track -------------//
  if(np1 == np0){
    for(int i = 0; i < NL; i++)
      if((*trk0)[i] && (*trk1)[i] && abs((*trk0)[i]->GetNu()-(*trk1)[i]->GetNu()) > VICINITY)
        // if(++vicinityViaCnt == np0) return true; // the two tracks are different
        if(++vicinityViaCnt == nCom) return true; // the two tracks are different
    // newtrk has better fitting quality to be accepted
    if(trk1->GetChi2() < trk0->GetChi2()){ fTracks[short(type)]->Remove(trk0); return true; }
    return false; // old track beats the new track in tracking quality
  } // end if(np1 > np0)

  return true; // never will reach here, just for suppressing compiler warnings
} // end member function Compare

void ETFTracker::Draw(){
  ConstructGraph();
  for(auto t : *fTracks[0]) static_cast<ETFVTrack *>(t)->Draw();
} // end member function Draw

void ETFTracker::ConstructGraph(){
  static const double L = ETFVisual::Instance()->GetLength();
  static const double W = ETFVisual::Instance()->GetWidth();
  for(const auto &t : *fTracks[0]) static_cast<ETFVTrack *>(t)->ConstructGraph(-W, L);
} // end member function ConstructGraph

void ETFTracker::Write(const char *name){
  if(!strcmp(name, "")) name = fTracks[0]->GetClass()->GetName();
  for(auto &c : fTracks) if(c && c->GetEntriesFast()) c->Write(name, TObject::kSingleKey);
  if(f3DTracks && f3DTracks->GetEntriesFast()) f3DTracks->Write(name, TObject::kSingleKey);
} // end member function Write

void ETFTracker::Print(){
  cout << "\033[32;1m=========== Tracker " << GetName();
  cout << " PRINT =============\033[0m" << endl;
  for(auto &c : fTracksOut) if(c && c->GetEntriesFast()){
    cout << "\033[33;1m_________ Number of " << c->GetName() << " tracks: ";
    cout << c->GetEntriesFast() << " __________\033[0m" << endl;
    for(auto p : *c) static_cast<ETFVTrack *>(p)->Print();
  } // end for
  if(f3DTracks){
    cout << "\033[47;35;1m______________ Number of " << f3DTracks->GetName() << " 3Dtracks: ";
    cout << f3DTracks->GetEntriesFast() << " _______________\033[0m" << endl;
    for(auto p : *f3DTracks) static_cast<ETF3DTrack *>(p)->Print();
  } // end if
} // end member function Print

// get the i-th track; t: 0-1-2: X-U(Y)-V
ETFVTrack *ETFTracker::GetTrack(int i, TrackType type){
  const short t = short(type);
  auto &c = fTracksOut[t];
  if(!c) er("ETFTracker", "GetTrack: fTracks[%d] is nullptr", t);
  if(i < 0 || i >= c->GetEntriesFast())
    er("ETFTracker", "GetTrack: i = %d out of range", i);

  int n = -1;
  // note that there may well be null elements between valid ones in a TClonesArray
  for(const auto &p : *c){
    if(p) n++;
    if(n == i){
      ETFVTrack *trk = static_cast<ETFVTrack *>(p);
      if(!trk) er("ETFTracker", "GetTrack: the requird trk is nullptr");
      return trk;
    } // end if
  } // end for

  return nullptr; // probably will never reach here
} // end member function GetTrack
// get the i-th 3D track
ETF3DTrack *ETFTracker::Get3DTrack(int i){
  auto &c = f3DTracks;
  if(!c) er("ETFTracker", "Get3DTrack: f3DTracks is nullptr");
  if(i < 0 || i >= c->GetEntriesFast())
    er("ETFTracker", "Get3DTrack: i = %d out of range", i);

  int n = -1;
  // note that there may well be null elements between valid ones in a TClonesArray
  for(const auto &p : *c){
    if(p) n++;
    if(n == i){
      ETF3DTrack *trk = static_cast<ETF3DTrack *>(p);
      if(!trk) er("ETFTracker", "Get3DTrack: the requird trk is nullptr");
      return trk;
    } // end if
  } // end for

  return nullptr; // probably will never reach here
} // end member function Get3DTrack

bool ETFTracker::SetDriftDistance(ETFVTrack *trk){
  for(auto &p : trk->GetHits()){
    ETFDCHit *h = p.second;
    if(!fDCs[h->GetLayerId()/2]->SetDriftDistance(h)) return false;
  } // end for
  return true;
} // end member function SetDriftDistance

///< if accept the layer or not in Map()
bool ETFTracker::GoodLayer(ETFVDC *dc, DCType type, int x12) const{
  if(!dc) er("ETFTracker", "GoodLayer: null dc pointer.");

  const int mul = dc->Multi(type, x12%2);
  // if(mul > 0.5*dc->GetNL()) return false;
  if(mul >= fMultiThrePerL) return false;
  return true;
} // end member function GoodLayer

ETFTimer *ETFTracker::GetTimer(int i) const{
  if(i >= 3 || !fTimer[i]) er("ETFTaTracker", "GetTimer: invalid input i = %d", i);
  return fTimer[i];
} // end member function GetTimer
